#coding:utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F


class STN(object):
    def __init__(self, input_shape, output_size=None, is_batch=True):
        """
        :input_shape not include batchsize just as (channel, height, width)
        :output_size is feature map's (height, width)
        """
        self.input_shape = input_shape
        self.output_size = output_size

        self.func = self.__batch_apply if is_batch else self.__apply

    def __batch_apply(self, batch_image, coef):
        # affine muset be reshape (-1, 2, 3)
        x = coef.view(-1, 2, 3)
        bs = x.size(0)
        if self.output_size:
            h, w = self.output_size
        else:
            h, w = self.input_shape[1:3]
        grid = F.affine_grid(x, torch.Size((bs, self.input_shape[0], h, w)))
        rois = F.grid_sample(batch_image, grid)
        return rois, grid

    def __apply(self, image, coef):
        x = coef.view(2, 3)
        grid = F.affine_grid(x, torch.Size(self.input_channel, self._h, self._w), align_corners=True)
        roi = F.grid_sample(image, grid, align_corners=True)
        return roi, grid
    
    def apply(self, x, coef):
        return self.func(x, coef) 


class STNLayer(STN):
    def __init__(self, *args, **kargs):
        super(STN, self).__init__(*args, **kargs)
        self.conv = nn.Conv2d(self.input_shape[0], 32, (3, 3), stride=1, padding=1, bias=False)
        self.lin1 = nn.Linear(32*self.input_shape[1] * self.input_shape[2], 1024)
        self.lin2 = nn.Linear(1024, 6)

    def forward(self, x):
        batch_image = x

        x = self.conv(x)
        x = F.relu(x)
        x = x.view(-1, 32*self.input_shape[1] * self.input_shape[2])
        x = self.lin1(x)
        x = self.lin2(x)

        return self.apply(batch_image, x)


if __name__ == '__main__':
    net = STNLayer(input_shape =(3,10,10))
    x = torch.randn((5, 3, 10, 10))
    y = net(x)
    print(y[0].size())